#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

// LLVM IR生成器结构体
typedef struct {
    FILE* output;
    int label_count;
    int temp_count;
    int indent_level;
} IRGenerator;

// 符号表结构（简化版）
typedef struct {
    char* name;
    bool is_global;
} Symbol;

// 抽象语法树节点类型
typedef enum {
    AST_PROGRAM,
    AST_FUNCTION,
    AST_ASSIGNMENT,
    AST_BINARY_OP,
    AST_IF_STATEMENT,
    AST_WHILE_LOOP,
    AST_FUNCTION_CALL,
    AST_RETURN,
    AST_IDENTIFIER,
    AST_INTEGER
} ASTNodeType;

// 抽象语法树节点
typedef struct ASTNode {
    ASTNodeType type;
    union {
        // 函数定义
        struct {
            char* name;
            char** params;
            int param_count;
            struct ASTNode* body;
        } function;
        
        // 赋值语句
        struct {
            char* target;
            struct ASTNode* value;
        } assignment;
        
        // 二元操作
        struct {
            char op;
            struct ASTNode* left;
            struct ASTNode* right;
        } binary_op;
        
        // if语句
        struct {
            struct ASTNode* condition;
            struct ASTNode* then_branch;
            struct ASTNode* else_branch;
        } if_stmt;
        
        // while循环
        struct {
            struct ASTNode* condition;
            struct ASTNode* body;
        } while_loop;
        
        // 函数调用
        struct {
            char* name;
            struct ASTNode** args;
            int arg_count;
        } function_call;
        
        // return语句
        struct {
            struct ASTNode* value;
        } return_stmt;
        
        // 标识符
        char* identifier;
        
        // 整数字面量
        int integer;
    } data;
} ASTNode;

// 初始化IR生成器
IRGenerator* init_ir_generator(const char* output_file) {
    IRGenerator* gen = malloc(sizeof(IRGenerator));
    gen->output = fopen(output_file, "w");
    gen->label_count = 0;
    gen->temp_count = 0;
    gen->indent_level = 0;
    return gen;
}

// 释放IR生成器资源
void free_ir_generator(IRGenerator* gen) {
    if (gen->output) fclose(gen->output);
    free(gen);
}

// 生成缩进
void emit_indent(IRGenerator* gen) {
    for (int i = 0; i < gen->indent_level; i++) {
        fprintf(gen->output, "  ");
    }
}

// 生成IR代码
void emit(IRGenerator* gen, const char* format, ...) {
    emit_indent(gen);
    va_list args;
    va_start(args, format);
    vfprintf(gen->output, format, args);
    va_end(args);
    fprintf(gen->output, "\n");
}

// 生成唯一标签
char* new_label(IRGenerator* gen, const char* prefix) {
    char* label = malloc(32);
    sprintf(label, "%s_%d", prefix, ++gen->label_count);
    return label;
}

// 生成唯一临时变量
char* new_temp(IRGenerator* gen) {
    char* temp = malloc(16);
    sprintf(temp, "%%temp%d", ++gen->temp_count);
    return temp;
}

// 生成函数头部
void generate_function_header(IRGenerator* gen, ASTNode* node) {
    char* func_name = node->data.function.name;
    
    // 生成参数列表
    char params[256] = "";
    for (int i = 0; i < node->data.function.param_count; i++) {
        strcat(params, "i32 ");
        strcat(params, node->data.function.params[i]);
        if (i < node->data.function.param_count - 1) {
            strcat(params, ", ");
        }
    }
    
    emit(gen, "define i32 @%s(%s) {", func_name, params);
    gen->indent_level++;
    
    // 为参数分配空间
    for (int i = 0; i < node->data.function.param_count; i++) {
        emit(gen, "%%%s.addr = alloca i32, align 4", node->data.function.params[i]);
        emit(gen, "store i32 %s, i32* %%%s.addr, align 4", 
             node->data.function.params[i], node->data.function.params[i]);
    }
}

// 生成表达式代码并返回结果变量
char* generate_expression(IRGenerator* gen, ASTNode* node);

// 生成赋值语句
void generate_assignment(IRGenerator* gen, ASTNode* node) {
    char* target = node->data.assignment.target;
    char* value = generate_expression(gen, node->data.assignment.value);
    
    // 这里简化处理，假设所有变量都是局部变量
    emit(gen, "store i32 %s, i32* %%%s.addr, align 4", value, target);
    free(value);
}

// 生成二元操作
char* generate_binary_op(IRGenerator* gen, ASTNode* node) {
    char* left = generate_expression(gen, node->data.binary_op.left);
    char* right = generate_expression(gen, node->data.binary_op.right);
    char* temp = new_temp(gen);
    
    switch (node->data.binary_op.op) {
        case '+':
            emit(gen, "%s = add nsw i32 %s, %s", temp, left, right);
            break;
        case '-':
            emit(gen, "%s = sub nsw i32 %s, %s", temp, left, right);
            break;
        case '*':
            emit(gen, "%s = mul nsw i32 %s, %s", temp, left, right);
            break;
        case '/':
            emit(gen, "%s = sdiv i32 %s, %s", temp, left, right);
            break;
        case '<':
            emit(gen, "%s = icmp slt i32 %s, %s", temp, left, right);
            break;
        case '>':
            emit(gen, "%s = icmp sgt i32 %s, %s", temp, left, right);
            break;
        case '=': // 处理 ==
            emit(gen, "%s = icmp eq i32 %s, %s", temp, left, right);
            break;
        default:
            fprintf(stderr, "Unsupported operator: %c\n", node->data.binary_op.op);
            exit(1);
    }
    
    free(left);
    free(right);
    return temp;
}

// 生成if语句
void generate_if_statement(IRGenerator* gen, ASTNode* node) {
    char* cond = generate_expression(gen, node->data.if_stmt.condition);
    char* then_label = new_label(gen, "if.then");
    char* else_label = new_label(gen, "if.else");
    char* end_label = new_label(gen, "if.end");
    
    emit(gen, "br i1 %s, label %%%s, label %%%s", cond, then_label, else_label);
    free(cond);
    
    // then分支
    emit(gen, "%s:", then_label);
    if (node->data.if_stmt.then_branch) {
        generate_expression(gen, node->data.if_stmt.then_branch);
    }
    emit(gen, "br label %%%s", end_label);
    
    // else分支
    emit(gen, "%s:", else_label);
    if (node->data.if_stmt.else_branch) {
        generate_expression(gen, node->data.if_stmt.else_branch);
    }
    emit(gen, "br label %%%s", end_label);
    
    // 结束标签
    emit(gen, "%s:", end_label);
    
    free(then_label);
    free(else_label);
    free(end_label);
}

// 生成while循环
void generate_while_loop(IRGenerator* gen, ASTNode* node) {
    char* cond_label = new_label(gen, "while.cond");
    char* body_label = new_label(gen, "while.body");
    char* end_label = new_label(gen, "while.end");
    
    emit(gen, "br label %%%s", cond_label);
    
    // 条件判断
    emit(gen, "%s:", cond_label);
    char* cond = generate_expression(gen, node->data.while_loop.condition);
    emit(gen, "br i1 %s, label %%%s, label %%%s", cond, body_label, end_label);
    free(cond);
    
    // 循环体
    emit(gen, "%s:", body_label);
    if (node->data.while_loop.body) {
        generate_expression(gen, node->data.while_loop.body);
    }
    emit(gen, "br label %%%s", cond_label);
    
    // 结束标签
    emit(gen, "%s:", end_label);
    
    free(cond_label);
    free(body_label);
    free(end_label);
}

// 生成函数调用
char* generate_function_call(IRGenerator* gen, ASTNode* node) {
    char* func_name = node->data.function_call.name;
    char* temp = new_temp(gen);
    char args[256] = "";
    
    // 生成参数列表
    for (int i = 0; i < node->data.function_call.arg_count; i++) {
        char* arg = generate_expression(gen, node->data.function_call.args[i]);
        strcat(args, "i32 ");
        strcat(args, arg);
        if (i < node->data.function_call.arg_count - 1) {
            strcat(args, ", ");
        }
        free(arg);
    }
    
    // 特殊处理printf
    if (strcmp(func_name, "printf") == 0) {
        emit(gen, "%s = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([15 x i8], [15 x i8]* @.str, i32 0, i32 0), i32 %s", 
             temp, args);
    } else {
        emit(gen, "%s = call i32 @%s(%s)", temp, func_name, args);
    }
    
    return temp;
}

// 生成return语句
void generate_return(IRGenerator* gen, ASTNode* node) {
    if (node->data.return_stmt.value) {
        char* value = generate_expression(gen, node->data.return_stmt.value);
        emit(gen, "ret i32 %s", value);
        free(value);
    } else {
        emit(gen, "ret i32 0");
    }
}

// 生成标识符
char* generate_identifier(IRGenerator* gen, ASTNode* node) {
    char* id = node->data.identifier;
    char* temp = new_temp(gen);
    
    // 简化处理：假设所有变量都是局部变量
    emit(gen, "%s = load i32, i32* %%%s.addr, align 4", temp, id);
    return temp;
}

// 生成整数字面量
char* generate_integer(IRGenerator* gen, ASTNode* node) {
    char* temp = malloc(16);
    sprintf(temp, "%d", node->data.integer);
    return temp;
}

// 生成表达式代码并返回结果变量
char* generate_expression(IRGenerator* gen, ASTNode* node) {
    switch (node->type) {
        case AST_ASSIGNMENT:
            generate_assignment(gen, node);
            return NULL;
        case AST_BINARY_OP:
            return generate_binary_op(gen, node);
        case AST_IF_STATEMENT:
            generate_if_statement(gen, node);
            return NULL;
        case AST_WHILE_LOOP:
            generate_while_loop(gen, node);
            return NULL;
        case AST_FUNCTION_CALL:
            return generate_function_call(gen, node);
        case AST_RETURN:
            generate_return(gen, node);
            return NULL;
        case AST_IDENTIFIER:
            return generate_identifier(gen, node);
        case AST_INTEGER:
            return generate_integer(gen, node);
        default:
            fprintf(stderr, "Unsupported AST node type: %d\n", node->type);
            exit(1);
    }
}

// 生成函数体
void generate_function_body(IRGenerator* gen, ASTNode* node) {
    // 递归生成函数体中的语句
    ASTNode* body = node->data.function.body;
    
    // 这里简化处理：假设函数体是一个语句列表
    // 实际实现中应处理语句序列
    generate_expression(gen, body);
    
    // 如果没有返回语句，添加默认返回
    if (body->type != AST_RETURN) {
        emit(gen, "ret i32 0");
    }
}

// 生成整个函数
void generate_function(IRGenerator* gen, ASTNode* node) {
    generate_function_header(gen, node);
    generate_function_body(gen, node);
    
    gen->indent_level--;
    emit(gen, "}\n");
}

// 生成模块头部
void generate_header(IRGenerator* gen) {
    emit(gen, "; ModuleID = 'example.c'");
    emit(gen, "source_filename = \"example.c\"\n");
    emit(gen, "@i = global i32 0, align 4");
    emit(gen, "@.str = private unnamed_addr constant [15 x i8] c\"Value: %d\\0A\\00\", align 1\n");
    emit(gen, "declare i32 @printf(i8*, ...)\n");
}

// 生成整个程序
void generate_program(IRGenerator* gen, ASTNode* node) {
    generate_header(gen);
    
    // 遍历所有顶层节点（函数）
    // 这里简化处理：假设节点是函数列表
    generate_function(gen, node);
}

// 创建AST节点函数
ASTNode* create_function_node(char* name, char** params, int param_count, ASTNode* body) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_FUNCTION;
    node->data.function.name = strdup(name);
    node->data.function.params = params;
    node->data.function.param_count = param_count;
    node->data.function.body = body;
    return node;
}

ASTNode* create_assignment_node(char* target, ASTNode* value) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_ASSIGNMENT;
    node->data.assignment.target = strdup(target);
    node->data.assignment.value = value;
    return node;
}

ASTNode* create_binary_op_node(char op, ASTNode* left, ASTNode* right) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_BINARY_OP;
    node->data.binary_op.op = op;
    node->data.binary_op.left = left;
    node->data.binary_op.right = right;
    return node;
}

ASTNode* create_if_node(ASTNode* condition, ASTNode* then_branch, ASTNode* else_branch) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_IF_STATEMENT;
    node->data.if_stmt.condition = condition;
    node->data.if_stmt.then_branch = then_branch;
    node->data.if_stmt.else_branch = else_branch;
    return node;
}

ASTNode* create_while_node(ASTNode* condition, ASTNode* body) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_WHILE_LOOP;
    node->data.while_loop.condition = condition;
    node->data.while_loop.body = body;
    return node;
}

ASTNode* create_function_call_node(char* name, ASTNode** args, int arg_count) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_FUNCTION_CALL;
    node->data.function_call.name = strdup(name);
    node->data.function_call.args = args;
    node->data.function_call.arg_count = arg_count;
    return node;
}

ASTNode* create_return_node(ASTNode* value) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_RETURN;
    node->data.return_stmt.value = value;
    return node;
}

ASTNode* create_identifier_node(char* name) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_IDENTIFIER;
    node->data.identifier = strdup(name);
    return node;
}

ASTNode* create_integer_node(int value) {
    ASTNode* node = malloc(sizeof(ASTNode));
    node->type = AST_INTEGER;
    node->data.integer = value;
    return node;
}

// 示例使用
int main() {
    // 创建示例程序AST：计算阶乘并打印结果
    // int main() {
    //   i = 5;
    //   printf(factorial(i));
    //   return 0;
    // }
    
    // 创建参数列表
    char* main_params[] = {};
    
    // 创建main函数体
    ASTNode* assign_i = create_assignment_node("i", create_integer_node(5));
    ASTNode* call_fact = create_function_call_node(
        "factorial", 
        (ASTNode*[]){create_identifier_node("i")}, 
        1
    );
    ASTNode* call_printf = create_function_call_node(
        "printf", 
        (ASTNode*[]){call_fact}, 
        1
    );
    ASTNode* return0 = create_return_node(create_integer_node(0));
    
    // 创建语句序列（简化处理）
    ASTNode* main_body = call_printf;
    
    ASTNode* main_func = create_function_node("main", main_params, 0, main_body);
    
    // 创建阶乘函数
    // int factorial(int n) {
    //   if (n <= 1) return 1;
    //   else return n * factorial(n-1);
    // }
    
    char* fact_params[] = {"n"};
    ASTNode* condition = create_binary_op_node(
        '<', 
        create_identifier_node("n"), 
        create_integer_node(2) // n <= 1 等价于 n < 2
    );
    
    ASTNode* then_branch = create_return_node(create_integer_node(1));
    
    ASTNode* recursive_call = create_function_call_node(
        "factorial", 
        (ASTNode*[]){
            create_binary_op_node(
                '-', 
                create_identifier_node("n"), 
                create_integer_node(1)
            )
        }, 
        1
    );
    
    ASTNode* else_branch = create_return_node(
        create_binary_op_node(
            '*', 
            create_identifier_node("n"), 
            recursive_call
        )
    );
    
    ASTNode* if_stmt = create_if_node(condition, then_branch, else_branch);
    ASTNode* fact_func = create_function_node("factorial", fact_params, 1, if_stmt);
    
    // 初始化IR生成器
    IRGenerator* gen = init_ir_generator("output.ll");
    
    // 生成阶乘函数
    generate_function(gen, fact_func);
    
    // 生成main函数
    generate_function(gen, main_func);
    
    // 释放资源
    free_ir_generator(gen);
    
    printf("LLVM IR generated in output.ll\n");
    return 0;
}